package comfy

import (
	"context"
	"errors"
	"fmt"
	"log"
	"os"
	"quipus/internal/defs"
	"quipus/internal/llm"
	comfyClient "quipus/internal/llm/provider/comfy/client"
)

type ComfyClient struct {
	client *comfyClient.Client
}

func NewClient(config *llm.LLMClientConfig) (*ComfyClient, error) {
	client, err := comfyClient.NewDefaultClientStr(config.BaseUrl)
	if err != nil {
		return nil, err
	}
	c := &ComfyClient{
		client: client,
	}

	client.ConnectAndListen()

	for !client.IsInitialized() {
	}

	return c, nil
}

func IsEndQueuePrompt(count int, num int) bool {
	if count >= num {
		fmt.Println("end queue prompt")
		return true
	}
	return false
}

func (c *ComfyClient) Image(ctx context.Context, req *llm.ImageReq) (rsp []llm.ImageInfo, err error) {
	if req == nil {
		return nil, nil
	}

	client := c.client

	// go func() {
	// 	_, err := client.QueuePromptByString(req.Prompt, req.Extra)
	// 	if err != nil {
	// 		fmt.Println(err)
	// 		os.Exit(1)
	// 	}
	// }()

	// go func() {
	// 	_, err := client.QueuePromptByNodes(getNodes(), req.Extra)
	// 	if err != nil {
	// 		fmt.Println(err)
	// 		os.Exit(1)
	// 	}
	// }()

	// ss := client.GetTaskStatus()
	// if len(ss) == 0 {
	// 	return "", nil
	// }

	promptRsp, e := client.QueuePromptByString(req.Prompt, req.Extra)
	if e != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	if len(promptRsp.NodeErrors) > 1 {
		return nil, errors.New("")
	}

	req.ID <- promptRsp.PromptID
	req.QueueCh <- int32(promptRsp.Number)

	count := 0
	for taskStatus := range client.GetTaskStatus() {
		switch taskStatus.Type {
		case comfyClient.ExecutionStart:
			req.CmdTypeCh <- defs.CmdStart
		case comfyClient.ExecutionCached:
			s := taskStatus.Data.(*comfyClient.WSMessageDataExecutionCached)
			fmt.Printf("Type: %v, Data:%+v\n", comfyClient.ExecutionCached, s)
		case comfyClient.Executing:
			s := taskStatus.Data.(*comfyClient.WSMessageDataExecuting)
			fmt.Printf("Type: %v, Data:%+v\n", comfyClient.Executing, s)

		case comfyClient.Progress:
			s := taskStatus.Data.(*comfyClient.WSMessageDataProgress)
			progress := llm.Progress{
				Current: s.Value,
				Max:     s.Max,
			}
			req.ProgressCh <- progress

		case comfyClient.ExecutionSuccess:
			s := taskStatus.Data.(*comfyClient.WSMessageExecutionSuccess)
			fmt.Printf("Type: %v, Data:%+v\n", comfyClient.Progress, s)
			return
		case comfyClient.Executed:
			s := taskStatus.Data.(*comfyClient.WSMessageDataExecuted)
			// fmt.Printf("Type: %v, Data:%+v\n", comfyClient.Executed, s)
			for _, images := range s.Output {
				for _, image := range images {
					imageData, err := client.GetFile(image)
					if err != nil {
						panic(err)
					}
					// f, err := os.Create(image.Filename)
					if err != nil {
						log.Println("Failed to write image:", err)
						os.Exit(1)
					}
					rsp = append(rsp, llm.ImageInfo{FileName: image.Filename, Img: *imageData})
					// f.Write(*imageData)
					// f.Close()
				}
			}
			count++
			if IsEndQueuePrompt(count, 2) {
				req.CmdTypeCh <- defs.CmdExecuted
			}
			return

		case comfyClient.ExecutionInterrupted:
			s := taskStatus.Data.(*comfyClient.WSMessageExecutionInterrupted)
			fmt.Printf("Type: %v, Data:%+v\n", comfyClient.ExecutionInterrupted, s)
			// count++
			// IsEndQueuePrompt(count, 2)
			err = errors.New("execution interrupted")
			return
		case comfyClient.ExecutionError:
			s := taskStatus.Data.(*comfyClient.WSMessageExecutionError)
			fmt.Printf("Type: %v, Data:%+v\n", comfyClient.ExecutionError, s)
			count++
			IsEndQueuePrompt(count, 2)
			err = errors.New(s.ExceptionMessage)
			return
		default:
			fmt.Println("unknown message type")
		}
	}

	return nil, nil

}
